import os
import json
import time
import logging
import collections
from pack_to_json.reader import Reader
import re

# 列举所需要的topic
ADAS_topic_list = ['runtime', 'log', 'online_calib', 'contour_endpts', 'can', 'lane_parsing', 'crop_roi', 'real3d',
                   'resize_roi', 'multi_cam_state', 'runtime_env',
                   'cam_freespace', 'skeleton_raw', 'boxes_3d_raw', 'obstacles_raw', 'traffic_light', 'traffic_sign',
                   'object_parsing_bbox','parsing_bbox_v2', 'parsing', 'lane_result', 'camera_default', 'object_pole',
                   'freespace_points', 'odometry', 'obstacle', 'obs_track_run', 'obs_state_run',
                   'postinput_runmsg', 'vehicle_result', 'ped_result', 'laneline_v2', 'object_light', 'object_sign',
                   'scan_points', 'image_quality', 'ihbc', 'weather_type', 'time_type', 'light_type',
                   'scene_type', 'image']
NDMLOC_topic_list = ['location']
NDMDEBUG_topic_list = ['loc_debug_msg']
NDMENV_topic_list = ['map_env_v2']
NDMREC_topic_list = ['rec_env_v2']
NDMNAV_topic_list = ['fus_nav_v2']
NDMFUS_topic_list = ['fus_env_v2']
Robo_topic_list = ['odometry', 'obs_track_run', 'obs_state_run', 'obstacle']
RoboRaw_topic_list = ['can', 'lidar_perception_raw', 'odometry', 'obs_track_run', 'obs_state_run', 'obstacle', 'robosense', 'runtime_env']
TRAJPRED_topic_list = ['run_log']
PNC_topic_list = ['ssm_road', 'planning_traj', 'control_cmd']
BEV_RAW_topic_list = ['runtime_env', 'bev_parsing', 'odometry', 'bev3d', 'multi_cam_state', 'online_mapping', 'bev_parsing_conf']
FUSION_topic_list = ['odometry', 'runtime', 'obs_track_run', 'obs_state_run', 'obstacle', 'postinput_runmsg', 'multi_cam_state']
Radar_topic_list = ['radar']
F9P_topic_list = ['f9p']
F9Kimu_topic_list = ['asm330']
HMI_topic_list = ['hmi_drive']

Pack_type_dict = {'ADAS': ADAS_topic_list,
                  'NDMLOC': NDMLOC_topic_list,
                  'NDMREC': NDMREC_topic_list,
                  'NDMENV': NDMENV_topic_list,
                  'NDMNAV': NDMNAV_topic_list,
                  'NDMFUS': NDMFUS_topic_list,
                  'NDMDEBUG': NDMDEBUG_topic_list,
                  'Robo': Robo_topic_list,
                  'RoboRaw': RoboRaw_topic_list,
                  'PREDICTION': TRAJPRED_topic_list,
                  'PNC': PNC_topic_list,
                  'FUSION6V': FUSION_topic_list,
                  'BEV-RAW': BEV_RAW_topic_list,
                  'BEV-POST-PROCESS': FUSION_topic_list,
                  'MULFUSION': FUSION_topic_list,
                  'REAR-POST-PROCESS': FUSION_topic_list,
                  'radar#': Radar_topic_list,
                  'F9P#': F9P_topic_list,
                  'F9Kimu#': F9Kimu_topic_list,
                  'HMI': HMI_topic_list
                  }


def dict2json(proto_dict, pack_name, json_path):
    print("****************dict2json*******************")
    json_str = json.dumps(proto_dict, indent=4)
    json_name = pack_name.replace(".pack", ".json")
    print('json_path:{}'.format(json_path))
    with open(json_path, 'w') as json_file:
        json_file.write(json_str)


# 递归实现proto转换成dic
def proto2dict(proto, dict, field):
    print('proto2dict:{}'.format(proto))
    if isinstance(proto, (str, int, float, bool)):
        dict[field] = proto
    elif isinstance(proto, bytes):
        dict[field] = str(proto)
    else:
        if hasattr(proto, '__len__'):
            field_list = []
            for sub_proto in proto:
                tmp_dict = collections.OrderedDict()
                proto2dict(sub_proto, tmp_dict, field)
                field_list.append(tmp_dict[field])
            dict[field] = field_list

        elif hasattr(proto, 'ListFields'):
            dict[field] = collections.OrderedDict()
            for fieldDescriptor in proto.ListFields():
                attr_name = fieldDescriptor[0].name
                attr_object = getattr(proto, attr_name)
                proto2dict(attr_object, dict[field], attr_name)

        else:
            print('proto2dict unknow type field:%s' % field)


def analysis_frame_v1(frame, frame_dict, topics_obj):
    data_descriptor = frame.frame_v1.data_descriptor
    image = frame.frame_v1.image[0]
    obstacles = frame.frame_v1.structure_perception.obstacles[0]
    lines = frame.frame_v1.structure_perception.lines[0]
    scan_pts = frame.frame_v1.structure_perception.scan_pts[0]

    proto2dict(obstacles, frame_dict, 'obstacles')
    proto2dict(lines, frame_dict, 'lines')
    proto2dict(scan_pts, frame_dict, 'scan_pts')
    proto2dict(image, frame_dict, 'image')

    for dd in data_descriptor:
        if dd.HasField("data") and dd.data.with_data_field == False:
            topic_name = dd.type
            if topic_name in ['image_quality', 'ihbc', 'weather_type', 'time_type', 'light_type', 'scene_type']:
                cur_obj = topics_obj[topic_name]()
                cur_obj.ParseFromString(dd.data.proto)
                proto2dict(cur_obj, frame_dict, topic_name)


def unpack(pack_name, topic_list, json_path):
    with open(pack_name, 'rb') as fd:
        reader = Reader(fd, topics=topic_list)
        topics_obj = reader.topics_obj
        frames_dict = collections.OrderedDict()

        for f in reader:
            frame = f.meta
            frame_id = frame.frame_id
            frames_dict[str(frame_id)] = collections.OrderedDict()

            if frame.proto_version == 2:
                topic_meta = f.topic_meta
                data_list = f.data_list
                for topic in topic_list:
                    if not (topic in topic_meta):
                        continue
                    print('topic:{}, topic_meta:{}'.format(topic, topic_meta))

                    topic_params = topic_meta[topic]
                    for i in range(len(topic_params)):
                        print('i:{}'.format(i))
                        if topic == "bev_parsing":
                            parsing_list = []
                            a = re.findall(r"\d+\.?\d*", str(data_list[i].read()))
                            for value in a:
                                if value not in parsing_list:
                                    parsing_list.append(value)
                            parsing_list_file = pack_name.rsplit("/", 1)[0] + "//bev_parsing.txt"
                            with open(parsing_list_file, 'a') as bev_file:
                                bev_file.write(str(parsing_list) + '\n')
                            bev_file.close()
                        if topic == "lane_parsing":
                            parsing_list = []
                            a = re.findall(r"\d+\.?\d*", str(data_list[i].read()))
                            for value in a:
                                if value not in parsing_list:
                                    parsing_list.append(value)
                            parsing_list_file = pack_name.rsplit("/", 1)[0] + "//front_parsing.txt"
                            with open(parsing_list_file, 'a') as front_file:
                                front_file.write(str(parsing_list) + '\n')
                            front_file.close()
                        if i == 0:
                            proto2dict(topic_params[i], frames_dict[str(frame_id)], topic)
                        else:
                            proto2dict(topic_params[i], frames_dict[str(frame_id)], topic + str(i))
            if frame.proto_version == 1:
                analysis_frame_v1(frame, frames_dict[str(frame_id)], topics_obj)
                # print(frames_dict[str(frame_id)])
            
            # if len(frames_dict) >= 1000:
            #     break

        dict2json(frames_dict, pack_name, json_path)
    print('end unpack pack:%s' % pack_name)


def pack2json(input_path, json_path, topic_list=None):
    print('input_path:', input_path)
    input_path = os.path.abspath(input_path)

    if not os.path.exists(input_path):
        print("pack2json path:%s not exist" % input_path)
        return False

    if os.path.isdir(input_path):
        print("pack2json path:%s not file" % input_path)
        return False

    pack_name = os.path.splitext(os.path.basename(input_path))
    print(pack_name)
    if pack_name[1] == '.pack':
        pack_type = pack_name[0].split('_')[0]
        print(pack_type)
        if pack_type in Pack_type_dict:
            if topic_list is None:
                unpack(input_path, Pack_type_dict[pack_type], json_path)
            else:
                unpack(input_path, topic_list, json_path)
        else:
            print("pack2json pack_type:%s not exist" % pack_type)
            return False

    return True


def pack_to_json(env_json):
    print("*************************start unpack*****************************")
    for filepath, dirnames, filenames in os.walk(env_json['save_path']):
        for filename in filenames:
            pack_full_name = os.path.join(filepath, filename)
            print(pack_full_name)
            if not pack2json(os.path.join(filepath, filename)):
                return -1
    return 1


if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s-%(levelname)s: %(message)s', level=logging.INFO)
    pack2json(r'C:\Users\minjie.hu\Desktop\SD_auto\saved_pack\1_Pedestrian\BEV_20220516-194207_843_7.pack')
